# -*- coding:utf-8 -*-

import os
import numpy as np
import torch
import torch.nn.parallel
import torch.optim
import nibabel as nib
from DataPreProcess import preprocess
from torch import nn
from torch.backends import cudnn


def main_test():
    ##############################
    # Using GPU
    ##############################
    # os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    # model = torch.load('model_epoch_10.pkl')
    # model = model.cuda()
    # model = nn.DataParallel(model)
    # cudnn.benchmark = True


    ##############################
    # Using CPU
    ##############################
    model = torch.load('model_epoch_10.pkl', map_location=lambda storage, loc: storage)
    model = model.module

    dataPath = os.path.abspath('.')+"/LCTSC-Test-S3-204"
    savePath =os.path.abspath('.')
    data = preprocess(dataPath)
    print("Finish pre-processing")

    predProbMapsPerClass = []
    # for i in range(len(data)):
    for i in range(100, 105): #for cpu test
        x = data[i]
        x = x[np.newaxis, :]
        output = test(x, model)
        output = output.data.numpy()
        predProbMapsPerClass.append(output)


    predProbMapsPerClass = np.array(predProbMapsPerClass)
    predProbMapsPerClass = predProbMapsPerClass.reshape(predProbMapsPerClass.shape[0], predProbMapsPerClass.shape[2], predProbMapsPerClass.shape[3], predProbMapsPerClass.shape[4])

    predSegmentation = np.argmax(predProbMapsPerClass, axis=1)

    seg = np.zeros(predSegmentation.shape, dtype=np.int16)
    seg[predSegmentation == 1] = 1
    seg[predSegmentation == 2] = 2
    seg[predSegmentation == 3] = 3
    seg[predSegmentation == 4] = 4
    seg[predSegmentation == 5] = 5

    img = nib.Nifti1Image(seg, np.eye(4))
    nib.save(img, savePath + '/test.nii.gz')



def test(x,model):

    model.eval()

    x =torch.from_numpy(np.array(x))
    x1 = torch.autograd.Variable(x)
    # x1 = torch.autograd.Variable(x.cuda())

    output = model(x1)  # nx5x9x9x9
    softmax = torch.nn.Softmax()
    output = output.view(-1, 6, 512 ** 2).permute(0, 2, 1).contiguous()
    output = output.view(-1, 6)

    output = softmax(output)
    output = output.view(-1, 512 ** 2, 6).permute(0, 2, 1).contiguous()
    output = output.view(-1, 6, 512, 512)
    return output


if __name__ == '__main__':
    main_test()

